# visual_verify_7scenes.py

import os
import os.path as osp
import sys
import glob
import logging
import argparse
import numpy as np
from tqdm import tqdm

# --- Import from our visual utility file ---
from visual_util import launch_visualizer

# --- Import from the original utility file for data loading helpers ---
# Note: Assumes verify_utils.py contains your updated imread_cv2 function
try:
    from verify_utils import setup_logging, depthmap_to_absolute_camera_coordinates, imread_cv2
except ImportError:
    print("FATAL ERROR: Could not import from 'verify_utils.py'.")
    print("Please ensure 'verify_utils.py' is in the same directory.")
    sys.exit(1)


def load_7scenes_view_data(rgb_path, intrinsics, logger):
    """
    Loads all necessary data for a single view from the 7-Scenes dataset format.
    """
    try:
        # Construct paths based on the 7-Scenes naming convention
        depth_path = rgb_path.replace("color.png", "depth.png")
        pose_path = rgb_path.replace("color.png", "pose.txt")

        if not all(osp.exists(p) for p in [depth_path, pose_path]):
            logger.warning(f"Missing files for {osp.basename(rgb_path)}. Skipping.")
            return None

        # Load image and depth data
        rgb_image = imread_cv2(rgb_path) # Assumes it returns RGB
        depthmap = imread_cv2(depth_path, -1) # Use -1 for cv2.IMREAD_UNCHANGED
        
        if rgb_image is None or depthmap is None:
            logger.warning(f"Failed to read image or depth for {rgb_path}. Skipping.")
            return None

        # --- Process data according to 7-Scenes format ---
        # 1. Normalize RGB image for the viewer
        rgb_image = rgb_image.astype(np.float32) / 255.0
        # 2. Handle invalid depth values
        depthmap[depthmap == 65535] = 0
        # 3. Scale depth from millimeters (uint16) to meters (float32)
        depthmap = depthmap.astype(np.float32) / 1000.0
        # 4. Handle any remaining non-finite values
        depthmap[~np.isfinite(depthmap)] = 0

        # Load camera pose
        camera_pose = np.loadtxt(pose_path, delimiter=None).astype(np.float32)

        return {
            "rgb": rgb_image, "depth": depthmap, "intrinsics": intrinsics,
            "pose": camera_pose, "path": rgb_path
        }
    except Exception as e:
        logger.error(f"Could not process view for {osp.basename(rgb_path)}. Details: {e}", exc_info=True)
        return None

# ======================================================================================
# --- Main Execution Block ---
# ======================================================================================

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Visually verify the consistency of a 7-Scenes dataset."
    )
    
    parser.add_argument(
        "--data_root", 
        type=str,
        default='/lc/data/3D/7scenes',
        help="Path to the root directory of the 7-Scenes dataset."
    )
    
    parser.add_argument(
        "--log_file", 
        type=str, 
        default="logs/visual_verification_7scenes.log",
        help="Path to the output log file."
    )

    parser.add_argument(
        "--scene", 
        type=str, 
        default='office',
        help="Optional: Specify a single scene name (e.g., 'chess') to visualize."
    )

    parser.add_argument(
        "--subscene", 
        type=str, 
        default='seq-03',
        help="Optional: Specify a single scene name (e.g., 'chess') to visualize."
    )

    parser.add_argument(
        "--port", 
        type=int, 
        default=8080,
        help="Port to run the Viser web server on."
    )
    
    parser.add_argument(
        "--use_kinectfusion_intrinsics",
        action="store_true",
        help="If set, use the KinectFusion intrinsics instead of the standard 7-Scenes ones."
    )
    
    args = parser.parse_args()
    # args.use_kinectfusion_intrinsics = True
    # Ensure the log directory exists
    log_dir = osp.dirname(args.log_file)
    if log_dir:
        os.makedirs(log_dir, exist_ok=True)
    logger = setup_logging(args.log_file)

    # --- Define Intrinsics based on user choice ---
    if args.use_kinectfusion_intrinsics:
        logger.info("Using KinectFusion intrinsics.")
        # Note: In this case, depth and color intrinsics are the same.
        depth_intrinsics = np.array([
            [585.0, 0.0, 320.0],
            [0.0, 585.0, 240.0],
            [0.0, 0.0, 1.0]
        ], dtype=np.float32)
    else:
        logger.info("Using standard 7-Scenes intrinsics.")
        # The point cloud is generated from the depth map, so we use depth intrinsics.
        depth_intrinsics = np.array([
            [598.84, 0.0, 320.0],
            [0.0, 587.62, 240.0],
            [0.0, 0.0, 1.0]
        ], dtype=np.float32)

    if not osp.isdir(args.data_root):
        logger.error(f"Provided data root '{args.data_root}' does not exist.")
        sys.exit(1)
    
    if args.scene:
        scene_dirs = [osp.join(args.data_root, args.scene)]
        if not osp.isdir(scene_dirs[0]):
            logger.error(f"Specified scene directory '{scene_dirs[0]}' does not exist.")
            sys.exit(1)
    else:
        scene_dirs = sorted([d for d in glob.glob(osp.join(args.data_root, '*')) if osp.isdir(d)])
    
    logger.info(f"Found {len(scene_dirs)} scenes to process in '{args.data_root}'.")
    
    for scene_dir in scene_dirs:
        if args.subscene:
            subscene_dirs = [osp.join(scene_dir, args.subscene)]
        else:
            subscene_dirs = sorted([d for d in glob.glob(osp.join(scene_dir, 'seq-*')) if osp.isdir(d)])
        for subscene_dir in subscene_dirs:
            scene_key = f"{osp.basename(scene_dir)}/{osp.basename(subscene_dir)}"
            logger.info(f"\n--- Processing subscene: {scene_key} ---")

            rgb_paths = sorted(glob.glob(osp.join(subscene_dir, "*color.png")))[:200]
            if not rgb_paths:
                logger.warning(f"No color images found in {subscene_dir}. Skipping.")
                continue

            pc_list, color_list, conf_list = [], [], []
            focals, pps, Rs, ts = [], [], [], []
            
            logger.info("Loading views and generating point clouds...")
            for rgb_path in tqdm(rgb_paths, file=sys.stdout):
                view_data = load_7scenes_view_data(rgb_path, depth_intrinsics, logger)
                if view_data:
                    world_points, valid_mask, _ = depthmap_to_absolute_camera_coordinates(
                        view_data["depth"], view_data["intrinsics"], view_data["pose"]
                    )
                    pc_list.append(world_points)
                    color_list.append(view_data['rgb'])
                    conf_list.append(valid_mask.astype(np.float32))

                    intrinsics = view_data['intrinsics']
                    pose = view_data['pose']
                    focals.append(intrinsics[0, 0])
                    pps.append((intrinsics[0, 2], intrinsics[1, 2]))
                    Rs.append(pose[:3, :3])
                    ts.append(pose[:3, 3])
        
            if not pc_list:
                logger.warning(f"No valid data could be loaded for {scene_key}. Skipping.")
                continue

            cam_dict = {'focal': focals, 'pp': pps, 'R': Rs, 't': ts}

            try:
                logger.info(f"Launching interactive visualizer for {scene_key} at http://localhost:{args.port}")
                logger.info("--> Close the browser tab or press Ctrl+C in the console to continue.")
                launch_visualizer(
                    pc_list=pc_list, color_list=color_list, conf_list=conf_list,
                    cam_dict=cam_dict, port=args.port
                )
            except KeyboardInterrupt:
                logger.info(f"\nVisualizer for {scene_key} closed, proceeding...")
            except Exception as e:
                logger.error(f"Failed to launch visualizer for {scene_key}: {e}", exc_info=True)
            break
